import time
import rtmplog
import hashlib
import os

class dstat():
    
    fv = 0          #第一个视频的时间戳
    fa = 0          #第一个音频的时间戳
    path = ""
    vcode = ""
    acode = ""
    ANUM = 2000    #该值记录最大数组值
    it = 0      #如果有1000个包的内容，则需要写入到文件中
    vtm = []
    vdat = []
    vcts = []
    adat = []
    atm = []
    acts = []
    d_num = 0
    vtype = []

    def __init__(self, app, stream):
        self.ANUM = 1000    #该值记录最大数组值
        self.vtm = []
        self.vdat = []
        self.adat = []
        self.atm = []
        self.d_num = 0
        str1 = ""+str(app)+"_"+str(stream)
        self.path = "/tmp/"+hashlib.md5(str1.encode(encoding='UTF-8')).hexdigest()
        if not os.path.exists(self.path):
            os.makedirs(self.path)
            
    def add_vdat(self, dat, cts):
        vt = int(time.time()*1000)
        if self.fv == 0:
            self.fv = vt
        self.vtm.append(vt)
        self.vdat.append(len(dat))
        self.vcts.append(cts)
        vtype = self.get_video_type(dat)
        self.vtype.append(vtype)
        
        self.it += 1
        if self.it > 100:
            self.write_local_file()
            self.it = 0

        if len(self.vtm) > self.ANUM:
            self.vtm = self.vtm[1:]
            self.vdat = self.vdat[1:]
            self.vcts = self.vcts[1:]
            self.vtype = self.vtype[1:]
        rtmplog.logger.info("接收到视频帧数:" +str(len(self.vtm)))

    def add_adat(self, dat, cts):
        at = int(time.time()*1000)
        if self.fa == 0:
            self.fa = at
        self.atm.append(at)
        self.adat.append(len(dat))
        self.acts.append(cts)
        self.get_audio_type(dat)
        if len(self.atm) > self.ANUM*2:
            self.atm = self.atm[1:]
            self.adat = self.adat[1:]
            self.acts = self.acts[1:]
        rtmplog.logger.info("接收到音频帧数:" +str(len(self.atm)))

    def get_a_stat():           #获取音频的统计数据
        ret = []
        for at, ad in atm, adat:
            tt = {}
            t1 = at - atm[0]
            tt[str(t1)] = ad
            ret.append(tt)
        
        return ret
    def get_v_stat(self):           #获取视频的统计数据
        ret = []
        for vt, vd in vtm, vdat:
            tt = {}
            t1 = vt - vtm[0]
            tt[str(t1)] = vd
            ret.append(tt)
        
        return ret

    def  get_video_type(self, dat):         #获取帧类型,视频帧有三种类型，I/B/P 以及每个帧的PTS
        #取dat的第一位进行分析：前四位是帧类型,后四位是编码类型；0001：关键帧 0010普通帧； 0111 = 7 是AVC编码H264
        a1 = int.from_bytes(dat[0:1], byteorder='big', signed=False)
        vtype = ""
        vf4 = a1>>4
        ve4 = a1&0x0F
        if vf4 == 1:
            vtype = 1
        elif vf4 == 2:
            vtype = 0
        else:
            vtype = 0
        
        if ve4 == 7 and self.vcode == "":
            self.vcode = "H264"      
        return vtype
    '''
    SoundFormat: 4 个 bit。取值 10，表示 AAC codec。
    SoundRate: 2 个 bit。取值 3，表示 44k 的采用频率。
    SoundSize: 1 个 bit。取值 1，表示音频样本用 16 bit 表示。
    SoundType: 1 个 bit。取值 1，表示立体声。
    '''

    def get_audio_type(self, dat):
        a1 = int.from_bytes(dat[0:1], byteorder='big', signed=False)
        acode = ""
        atype = ""
        a03 = a1>>4
        a45 = a1&0x0F>>2
        a6 =  a1&0x02
        a7 =  a1&0x01
        if a03 == 10 and self.acode == "":
            acode = acode + " AAC "
        if a45 ==  3 and self.acode == "":
            acode = acode + " 44k "
        if a6 == 1  and self.acode == "":
            acode = acode + " 16 bit "
        if a7 == 1  and self.acode == "":
            acode = acode + " stero "

        if self.acode == "":
            self.acode = acode
        return
    
    def write_local_file(self):   #每5秒钟左右，写入本地文件

        with open(self.path+"/rtmp.inf", 'w') as f:
            f.write("video_begin:"+str(self.fv)+"\n")
            f.write("audio_begin:"+str(self.fa)+"\n")
            f.write("video_code:"+self.vcode+"\n")
            f.write("audio_code:"+self.acode+"\n")
        with open(self.path+"/rtmp.v", 'w') as f:
            f.write(str(self.vtm)+"\n")
            f.write(str(self.vdat)+"\n")
            f.write(str(self.vcts)+"\n")
        with open(self.path+"/rtmp.a", 'w') as f:
            f.write(str(self.atm)+"\n")
            f.write(str(self.adat)+"\n")
            f.write(str(self.acts)+"\n")
        with open(self.path+"/rtmp.k", 'w') as f:
            f.write(str(self.vtm)+"\n")
            f.write(str(self.vtype)+"\n")
            f.write(str(self.vcts)+"\n")
